"""
At the command line, only need to run once to install the package via pip:

$ pip install --upgrade google-generativeai
"""
import os
import pathlib
import time
import json
import google.generativeai as genai
from semantic_nav.text_process import extract_json_code_blocks
import socket
from semantic_nav.log_utils import get_logger
logger = get_logger()
package_name = 'semantic_nav'


def extract_semantic_map_from_video(video_file_path: str):
    socket.setdefaulttimeout(180)
    GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
    if GOOGLE_API_KEY is None:
        raise ValueError("GOOGLE_API_KEY environment variable must be set.")
    genai.configure(api_key=GOOGLE_API_KEY, transport="rest")
    logger.info("Uploading file...")
    video_file = genai.upload_file(path=video_file_path, display_name="test video")
    logger.info(f"Completed upload: {video_file.uri}")
    while video_file.state.name == "PROCESSING":
        logger.info("Waiting for video to be processed.")
        time.sleep(10)
        video_file = genai.get_file(video_file.name)

    if video_file.state.name == "FAILED":
        raise ValueError(video_file.state.name)
    logger.info("Video processing complete: " + video_file.uri)

    instruction = """Suppose you are a robot equiped with a camera navigating in the world. You can see the video recorded by the camera and output the timestamp when you arrive at and leave each scene. Extract all the meaningful objects you think you can use as a waypoint for navigation in each scene. 
    Output as following json format:
    ```json
    [
        {
            "scene": scene_name,
            "arrive_time": arrive_time,
            "leave_time": leave_time,
            "waypoints": {
                object_name: closest_time,
                ...
            }
        },
        ...
    ]
    ```
    all time must be 'minutes:seconds' format or N/A meaning throughout the video.
    """
    model = genai.GenerativeModel(
        model_name="models/gemini-1.5-pro-latest", system_instruction=instruction
    )

    chat_session = model.start_chat(history=[])
    response = chat_session.send_message(video_file)

    # Open the file and write the Python object to the file
    json_text = extract_json_code_blocks(response.text)
    if len(json_text) > 0:
        json_dict = json.loads(json_text[0])
    logger.info("\ntime map as follow:")
    logger.info(json_dict)
    genai.delete_file(video_file.name)
    logger.info(f"Deleted file {video_file.uri}")

    return json_dict


def main(args=None):
    from ament_index_python.packages import get_package_prefix
    package_path = get_package_prefix(package_name)
    video_file_path = os.path.join(package_path, 'video/output_rgb.mp4')
    time_json_path = os.path.join(package_path, 'json/time.json')
    json_dict = extract_semantic_map_from_video(video_file_path)
    with open(str(time_json_path), "w") as f:
        json.dump(obj=json_dict, fp=f)


if __name__ == '__main__':
    main()
